
#include "HookFileOpt.h"


extern "C"
NTKERNELAPI 
NTSTATUS
NtCreateFile(
			 OUT PHANDLE             FileHandle,
			 IN ACCESS_MASK          DesiredAccess,
			 IN POBJECT_ATTRIBUTES   ObjectAttributes,
			 OUT PIO_STATUS_BLOCK    IoStatusBlock,
			 IN PLARGE_INTEGER       AllocationSize OPTIONAL,
			 IN ULONG                FileAttributes,
			 IN ULONG                ShareAccess,
			 IN ULONG                CreateDisposition,
			 IN ULONG                CreateOptions,
			 IN PVOID                EaBuffer OPTIONAL,
			 IN ULONG                EaLength );

typedef NTSTATUS (* FnNtCreateFile) (
									 OUT PHANDLE             FileHandle,
									 IN ACCESS_MASK          DesiredAccess,
									 IN POBJECT_ATTRIBUTES   ObjectAttributes,
									 OUT PIO_STATUS_BLOCK    IoStatusBlock,
									 IN PLARGE_INTEGER       AllocationSize OPTIONAL,
									 IN ULONG                FileAttributes,
									 IN ULONG                ShareAccess,
									 IN ULONG                CreateDisposition,
									 IN ULONG                CreateOptions,
									 IN PVOID                EaBuffer OPTIONAL,
									 IN ULONG                EaLength );

FnNtCreateFile OriginNtCreateFile = NULL;


NTSTATUS HookNtCreateFile(
						  OUT PHANDLE             FileHandle,
						  IN ACCESS_MASK          DesiredAccess,
						  IN POBJECT_ATTRIBUTES   ObjectAttributes,
						  OUT PIO_STATUS_BLOCK    IoStatusBlock,
						  IN PLARGE_INTEGER       AllocationSize OPTIONAL,
						  IN ULONG                FileAttributes,
						  IN ULONG                ShareAccess,
						  IN ULONG                CreateDisposition,
						  IN ULONG                CreateOptions,
						  IN PVOID                EaBuffer OPTIONAL,
						  IN ULONG                EaLength )
{
	KdPrint(("HookNtCreateFile \n"));
	if (ObjectAttributes != NULL && ObjectAttributes->ObjectName != NULL)
		KdPrint(("ObjectAttributes->ObjectName = %wZ \n", ObjectAttributes->ObjectName));
	if (OriginNtCreateFile)
		return OriginNtCreateFile(FileHandle, DesiredAccess, ObjectAttributes, IoStatusBlock, AllocationSize, FileAttributes,
		ShareAccess, CreateDisposition, CreateOptions, EaBuffer, EaLength);
	else
		return STATUS_UNSUCCESSFUL;
}



extern "C" NTKERNELAPI NTSTATUS NtDeleteFile( IN POBJECT_ATTRIBUTES ObjectAttributes );

typedef NTSTATUS (* FnNtDeleteFile)( IN POBJECT_ATTRIBUTES ObjectAttributes );

FnNtDeleteFile OriginNtDeleteFile = NULL;

NTSTATUS HookNtDeleteFile( IN POBJECT_ATTRIBUTES ObjectAttributes )
{
	KdPrint(("HookNtDeleteFile "));
	if (ObjectAttributes != NULL)
	{
		KdPrint(("%wZ", ObjectAttributes->ObjectName));
	}
	if (OriginNtDeleteFile)
		return OriginNtDeleteFile(ObjectAttributes);
	else
		return STATUS_UNSUCCESSFUL;
}

extern "C"
NTSTATUS ObQueryNameString(
					   IN      PVOID                    Object,
					   OUT POBJECT_NAME_INFORMATION ObjectNameInfo OPTIONAL,
					   IN      ULONG                    Length,
					   OUT     PULONG                   ReturnLength
						   );

typedef struct _GET_NAME_CONTROL {

	PCHAR allocatedBuffer;
	CHAR smallBuffer[256];

} GET_NAME_CONTROL, *PGET_NAME_CONTROL;

PUNICODE_STRING
SfGetFileName(
    IN PFILE_OBJECT FileObject,
    IN NTSTATUS CreateStatus,
    IN OUT PGET_NAME_CONTROL NameControl
    )
/*++
 
Routine Description:
 
    This routine will try and get the name of the given file object.  This
    is guaranteed to always return a printable string (though it may be NULL).
    This will allocate a buffer if it needs to.
 
Arguments:
    FileObject - the file object we want the name for
 
    CreateStatus - status of the create operation
 
    NameControl - control structure used for retrieving the name.  It keeps
        track if a buffer was allocated or if we are using the internal
        buffer.
 
Return Value:
 
    Pointer to the unicode string with the name
 
--*/
{
    POBJECT_NAME_INFORMATION nameInfo;
    NTSTATUS status;
    ULONG size;
    ULONG bufferSize;
 
    //
    //  Mark we have not allocated the buffer
    //
 
    NameControl->allocatedBuffer = NULL;
 
    //
    //  Use the small buffer in the structure (that will handle most cases)
    //  for the name
    //
 
    nameInfo = (POBJECT_NAME_INFORMATION)NameControl->smallBuffer;
    bufferSize = sizeof(NameControl->smallBuffer);
 
    //
    //  If the open succeeded, get the name of the file, if it
    //  failed, get the name of the device.
    //
         
    status = ObQueryNameString(
                  (NT_SUCCESS( CreateStatus ) ?
                    (PVOID)FileObject :
                    (PVOID)FileObject->DeviceObject),
                  nameInfo,
                  bufferSize,
                  &size );
 
    //
    //  See if the buffer was to small
    //
 
    if (status == STATUS_BUFFER_OVERFLOW) {
 
        //
        //  The buffer was too small, allocate one big enough
        //
 
        bufferSize = size + sizeof(WCHAR);
 
        NameControl->allocatedBuffer = (PCHAR) ExAllocatePoolWithTag( 
                                            NonPagedPool,
                                            bufferSize,
                                            '1234' );
 
        if (NULL == NameControl->allocatedBuffer) {
 
            //
            //  Failed allocating a buffer, return an empty string for the name
            //
 
            RtlInitEmptyUnicodeString(
                (PUNICODE_STRING)&NameControl->smallBuffer,
                (PWCHAR)(NameControl->smallBuffer + sizeof(UNICODE_STRING)),
                (USHORT)(sizeof(NameControl->smallBuffer) - sizeof(UNICODE_STRING)) );
 
            return (PUNICODE_STRING)&NameControl->smallBuffer;
        }
 
        //
        //  Set the allocated buffer and get the name again
        //
 
        nameInfo = (POBJECT_NAME_INFORMATION)NameControl->allocatedBuffer;
 
        status = ObQueryNameString(
                      FileObject,
                      nameInfo,
                      bufferSize,
                      &size );
    }
 
    //
    //  If we got a name and an error opening the file then we
    //  just received the device name.  Grab the rest of the name
    //  from the FileObject (note that this can only be done if being called
    //  from Create).  This only happens if we got an error back from the
    //  create.
    //
 
    if (NT_SUCCESS( status ) && 
                    !NT_SUCCESS( CreateStatus )) {
 
        ULONG newSize;
        PCHAR newBuffer;
        POBJECT_NAME_INFORMATION newNameInfo;
 
        //
        //  Calculate the size of the buffer we will need to hold
        //  the combined names
        //
 
        newSize = size + FileObject->FileName.Length;
 
        //
        //  If there is a related file object add in the length
        //  of that plus space for a separator
        //
 
        if (NULL != FileObject->RelatedFileObject) {
 
            newSize += FileObject->RelatedFileObject->FileName.Length + 
                       sizeof(WCHAR);
        }
 
        //
        //  See if it will fit in the existing buffer
        //
 
        if (newSize > bufferSize) {
 
            //
            //  It does not fit, allocate a bigger buffer
            //
 
            newBuffer = (PCHAR) ExAllocatePoolWithTag( 
                                    NonPagedPool,
                                    newSize,
                                    '1234' );
 
            if (NULL == newBuffer) {
 
                //
                //  Failed allocating a buffer, return an empty string for the name
                //
 
                RtlInitEmptyUnicodeString(
                    (PUNICODE_STRING)&NameControl->smallBuffer,
                    (PWCHAR)(NameControl->smallBuffer + sizeof(UNICODE_STRING)),
                    (USHORT)(sizeof(NameControl->smallBuffer) - sizeof(UNICODE_STRING)) );
 
                return (PUNICODE_STRING)&NameControl->smallBuffer;
            }
 
            //
            //  Now initialize the new buffer with the information
            //  from the old buffer.
            //
 
            newNameInfo = (POBJECT_NAME_INFORMATION)newBuffer;
 
            RtlInitEmptyUnicodeString(
                &newNameInfo->Name,
                (PWCHAR)(newBuffer + sizeof(OBJECT_NAME_INFORMATION)),
                (USHORT)(newSize - sizeof(OBJECT_NAME_INFORMATION)) );
 
            RtlCopyUnicodeString( &newNameInfo->Name, 
                                  &nameInfo->Name );
 
            //
            //  Free the old allocated buffer (if there is one)
            //  and save off the new allocated buffer address.  It
            //  would be very rare that we should have to free the
            //  old buffer because device names should always fit
            //  inside it.
            //
 
            if (NULL != NameControl->allocatedBuffer) {
 
                ExFreePool( NameControl->allocatedBuffer );
            }
 
            //
            //  Readjust our pointers
            //
 
            NameControl->allocatedBuffer = newBuffer;
            bufferSize = newSize;
            nameInfo = newNameInfo;
 
        } else {
 
            //
            //  The MaximumLength was set by ObQueryNameString to
            //  one char larger then the length.  Set it to the
            //  true size of the buffer (so we can append the names)
            //
 
            nameInfo->Name.MaximumLength = (USHORT)(bufferSize - 
                                  sizeof(OBJECT_NAME_INFORMATION));
        }
 
        //
        //  If there is a related file object, append that name
        //  first onto the device object along with a separator
        //  character
        //
 
        if (NULL != FileObject->RelatedFileObject) {
 
            RtlAppendUnicodeStringToString(
                    &nameInfo->Name,
                    &FileObject->RelatedFileObject->FileName );
 
            RtlAppendUnicodeToString( &nameInfo->Name, L"\\" );
        }
 
        //
        //  Append the name from the file object
        //
 
        RtlAppendUnicodeStringToString(
                &nameInfo->Name,
                &FileObject->FileName );
 
        ASSERT(nameInfo->Name.Length <= nameInfo->Name.MaximumLength);
    }
 
    //
    //  Return the name
    //
 
    return &nameInfo->Name;
}

extern "C"
NTKERNELAPI 
NTSTATUS
NtReadFile(IN HANDLE               FileHandle,
		   IN HANDLE               Event OPTIONAL,
		   IN PIO_APC_ROUTINE      ApcRoutine OPTIONAL,
		   IN PVOID                ApcContext OPTIONAL,
		   OUT PIO_STATUS_BLOCK    IoStatusBlock,
		   OUT PVOID               Buffer,
		   IN ULONG                Length,
		   IN PLARGE_INTEGER       ByteOffset OPTIONAL,
		   IN PULONG               Key OPTIONAL );


typedef NTSTATUS (* FnNtReadFile) (IN HANDLE               FileHandle,
		   IN HANDLE               Event OPTIONAL,
		   IN PIO_APC_ROUTINE      ApcRoutine OPTIONAL,
		   IN PVOID                ApcContext OPTIONAL,
		   OUT PIO_STATUS_BLOCK    IoStatusBlock,
		   OUT PVOID               Buffer,
		   IN ULONG                Length,
		   IN PLARGE_INTEGER       ByteOffset OPTIONAL,
		   IN PULONG               Key OPTIONAL );

FnNtReadFile OriginNtReadFile = NULL;

NTSTATUS HookNtReadFile(IN HANDLE     FileHandle,
		   IN HANDLE               Event OPTIONAL,
		   IN PIO_APC_ROUTINE      ApcRoutine OPTIONAL,
		   IN PVOID                ApcContext OPTIONAL,
		   OUT PIO_STATUS_BLOCK    IoStatusBlock,
		   OUT PVOID               Buffer,
		   IN ULONG                Length,
		   IN PLARGE_INTEGER       ByteOffset OPTIONAL,
		   IN PULONG               Key OPTIONAL )
{
	
	PFILE_OBJECT fileObject = NULL;
	ObReferenceObjectByHandle(FileHandle, 0, *IoFileObjectType, KernelMode, (PVOID *)&fileObject, NULL);
	if ((fileObject != NULL) && (fileObject->Vpb != NULL) && (fileObject->Vpb->RealDevice != NULL))
	{
		KdPrint(("ReadFile FileName = %wZ \n", &fileObject->FileName));
	}
	//else
	//	KdPrint(("HookNtReadFile \n"));

	ObDereferenceObject(fileObject);

	if (OriginNtReadFile)
		return OriginNtReadFile(FileHandle, Event, ApcRoutine, ApcContext, IoStatusBlock, Buffer, Length, ByteOffset, Key);
	else
		return	STATUS_ACCESS_DENIED;
}


BOOLEAN HookFile()
{
	KdPrint(("SetHook1 \n"));
	BOOLEAN b = SetHook(NtReadFile, HookNtReadFile, (PVOID*)&OriginNtReadFile);
	KdPrint(("SetHook = %d", b));	
	return TRUE;
}

VOID unHookFile()
{
	KdPrint(("DeleteHook1 \n"));
	UnHook((PVOID) NtReadFile);
}